from openai import OpenAI
from common import get_logger, safe_filename, session
import time
from config import OPENAI_BASE_URL, OPENAI_API_KEY, IMAGE_PRE_SAVE_DIR, IMAGE_SAVE_DIR
import os

from watermark import remove_watermark

logger = get_logger(__name__)


def generate_images_by_text(title, intro):
    client = OpenAI(base_url=OPENAI_BASE_URL, api_key=OPENAI_API_KEY)
    # 如果目录不存在则创建
    if not os.path.exists(IMAGE_PRE_SAVE_DIR):
        os.makedirs(IMAGE_PRE_SAVE_DIR)
    img_paths = []
    filenames = []
    for i in range(3):
        response = client.images.generate(
            model="doubao-seedream-3-0-t2i-250415",
            prompt=intro,
            size="1024x1024",
            response_format="url",
        )
        url = response.data[0].url
        logger.info(f"生成第{i+1}张图片成功:{url}")
        filename = safe_filename(i + 1, title, ".jpg")
        filenames.append(filename)
        img_path = os.path.join(IMAGE_PRE_SAVE_DIR, filename)
        with open(img_path, "wb") as file:
            file.write(session.get(url, timeout=10).content)
            img_paths.append(img_path)
        time.sleep(1)
    return {
        "img_paths": img_paths,
        "filenames": filenames
    }

def remove_images_watermark(filenames):
    if not os.path.exists(IMAGE_SAVE_DIR):
        os.makedirs(IMAGE_SAVE_DIR)
    image_paths = []
    for filename in filenames:
        input_path = os.path.join(IMAGE_PRE_SAVE_DIR, filename)
        out_path = os.path.join(IMAGE_SAVE_DIR, filename)
        remove_watermark(input_path, out_path)
        image_paths.append(out_path)
    return image_paths
